from collections import defaultdict
from datetime import datetime
from typing import Dict, List

from emi.bar.data import BarData


class SameTimeBarList:
    """
    相同时间的BarData归类。
    """
    def __init__(self,is_daily:bool):
        self._is_daily = is_daily
        self._same_time_map = defaultdict(list)

    def is_daily(self):
        """
        是否日频级别数据
        :return:
        """
        return self._is_daily

    def time_to_key(self, time:datetime):
        if self._is_daily:
            return time.strftime("%Y-%m-%d")
        else:
            return time.strftime("%Y-%m-%d %H:%M")

    def key_to_time(self, key: str) -> datetime:
        if self._is_daily:
            return datetime.strptime(key, "%Y-%m-%d")
        else:
            return datetime.strptime(key, "%Y-%m-%d %H:%M")

    def put(self,bar:BarData):
        _key = self.time_to_key(bar.datetime)
        self._same_time_map[_key].append(bar)

    def extend(self,bars:List[BarData]):
        for bar in bars:
            self.put(bar)

    def size(self):
        """
        返回归类的个数。
        :return:
        """
        return len(self._same_time_map)

    def to_map(self)->Dict[datetime,List[BarData]]:
        _data_map = {}
        for key,bar_list in self._same_time_map.items():
            _time = self.key_to_time(key)
            _data_map[_time] = bar_list
        return _data_map